package main

import "container/heap"

type PointHeap [][]int

func (h *PointHeap) Len() int {
	return len(*h)
}

func (h *PointHeap) Less(i, j int) bool {
	// x, y, dist
	return (*h)[i][2] < (*h)[j][2]
}

func (h *PointHeap) Swap(i, j int) {
	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
}

func (h *PointHeap) Push(x any) {
	(*h) = append((*h), x.([]int))
}
func (h *PointHeap) Pop() any {
	x := (*h)[len(*h)-1]
	*h = (*h)[:len(*h)-1]
	return x
}

func kClosest(points [][]int, k int) [][]int {
	h := &PointHeap{}
	for i, p := range points {
		dis := p[0]*p[0] + p[1]*p[1]
		points[i] = append(points[i], dis)
		*h = append(*h, points[i])
	}
	heap.Init(h)
	result := make([][]int, 0, k)
	for i := 0; i < k; i++ {
		x := heap.Pop(h).([]int)
		x = x[:2] // 删掉多余的距离
		result = append(result, x)
	}
	return result
}

func main() {
	res := kClosest([][]int{
		{3, 3},
		{5, -1},
		{-2, 4},
	}, 2)
	println(res)
}
